from zhipuai import ZhipuAI
from tool import search_papers, query_by_paper_id, query_by_title, query_by_title_contain, query_by_chunk_contain
import json

def parse_function_call(model_response, messages):
    # 处理函数调用结果，根据模型返回参数，调用对应的函数。
    # 调用函数返回结果后构造tool message，再次调用模型，将函数结果输入模型
    # 模型会将函数调用结果以自然语言格式返回给用户。
    if model_response.choices[0].message.tool_calls:
        tool_call = model_response.choices[0].message.tool_calls[0]
        args = tool_call.function.arguments
        function_result = {}
        if tool_call.function.name == "search_papers":
            function_result = search_papers(**json.loads(args))
        if tool_call.function.name == "query_by_paper_id":
            function_result = query_by_paper_id(**json.loads(args))
        if tool_call.function.name == "query_by_title":
            function_result = query_by_title(**json.loads(args))
        if tool_call.function.name == "query_by_title_contain":
            function_result = query_by_title_contain(**json.loads(args))
        if tool_call.function.name == "query_by_chunk_contain":
            function_result = query_by_chunk_contain(**json.loads(args))
        messages.append({
            "role": "tool",
            "content": f"{json.dumps(function_result)}",
            "tool_call_id":tool_call.id
        })
        response = client.chat.completions.create(
            model="glm-4-plus",
            messages=messages,
            tools=tools,
        )
        print(response.choices[0].message)
        messages.append(response.choices[0].message.model_dump())

client = ZhipuAI()
tools = [
    {
        "type": "function",
        "function": {
            "name": "search_papers",
            "description": "根据文本查询搜索论文片段。逻辑上是查找和 query 内容相似/相匹配的论文片段",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "description": "查询文本",
                        "type": "string"
                    },
                    "top_k": {
                        "description": "返回前top_k个结果",
                        "type": "int",
                    }
                },
                "required": ["query"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_by_paper_id",
            "description": "根据论文 ID 查询论文片段。需要精确匹配。",
            "parameters": {
                "type": "object",
                "properties": {
                    "paper_id": {
                        "description": "论文 ID",
                        "type": "str"
                    },
                    "top_k": {
                        "description": "返回前top_k个结果",
                        "type": "int",
                    }
                },
                "required": ["paper_id"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_by_title",
            "description": "根据论文标题查询论文片段。需要精确匹配。",
            "parameters": {
                "type": "object",
                "properties": {
                    "title": {
                        "description": "论文标题",
                        "type": "string"
                    },
                    "top_k": {
                        "description": "返回前top_k个结果",
                        "type": "int",
                    }
                },
                "required": ["title"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_by_title_contain",
            "description": "搜索标题中包含特定文本的论文片段。",
            "parameters": {
                "type": "object",
                "properties": {
                    "title": {
                        "description": "特定文本",
                        "type": "string"
                    },
                    "top_k": {
                        "description": "返回前top_k个结果",
                        "type": "int",
                    }
                },
                "required": ["title"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_by_chunk_contain",
            "description": "搜索论文内容片段中包含特定文本的片段。",
            "parameters": {
                "type": "object",
                "properties": {
                    "chunk": {
                        "description": "论文内容片段",
                        "type": "string"
                    },
                    "top_k": {
                        "description": "返回前top_k个结果",
                        "type": "int",
                    }
                },
                "required": ["chunk"]
            },
        }
    }
]

messages = []
messages.append({"role": "system", "content": "不要假设或猜测传入函数的参数值。如果用户的描述不明确，请要求用户提供必要信息"})
messages.append({"role": "user", "content": "帮我查询标题中含有llm的5篇论文"})
response = client.chat.completions.create(
    model="glm-4-plus",
    messages=messages,
    tools=tools,
)
print(response.choices[0].message)
messages.append(response.choices[0].message.model_dump())
parse_function_call(response,messages)